#!/usr/bin/python3
# coding=utf-8

# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import os
import sys
import logging

import numpy as np

IS_OUTPUT_TXT = False


class QuantMatmulGenData:
    def __init__(self, m, n, k, quant_mode, is_trans_a, is_trans_b, is_bias, data_type_str):
        self.m = m
        self.n = n
        self.k = k
        self.quant_mode = quant_mode
        self.is_bias = is_bias
        self.is_trans_a = is_trans_a
        self.is_trans_b = is_trans_b
        self.data_type_str = data_type_str


    def gen_golden_data_int8_float16_dequant(self, work_dir):
        src_type = np.int8
        dst_type = np.float16
        data_type = np.int32
        quant_type = np.uint64
        temp_type = np.float32

        # generate input x1, x2
        x1_gm = np.random.randint(-10, 10, [self.m, self.k]).astype(src_type)
        x2_gm = np.random.randint(-10, 10, [self.k, self.n]).astype(src_type)
        if self.is_bias:
            bias_gm = np.random.randint(-2, 2, [1, self.n]).astype(data_type)
        if self.quant_mode == 2: # 1: scalar quant mode, 2: vector quant mode 
            quant_vector = np.random.uniform(0.1, 2.0, [1, self.n]).astype(temp_type)
            quant_vector_gm = np.frombuffer(quant_vector, data_type)
            quant_vector_gm = quant_vector_gm.astype(quant_type)
            
        y_gm_int32 = np.matmul(x1_gm.astype(data_type), x2_gm.astype(data_type))
        if self.is_bias:
            y_gm_int32 = y_gm_int32 + bias_gm

        golden = y_gm_int32.astype(dst_type)
        if self.quant_mode == 1:
            golden = golden * 0.1
        elif self.quant_mode == 2:
            quant_vector = quant_vector.view("uint32")
            for index, data in enumerate(quant_vector):
                # 1 sign bit, 8 exponent bits and 10 mantissa bits
                quant_vector[index] = np.bitwise_and(data, 0xFFFFE000)
            quant_vector = quant_vector.view("float32")
            for i in range(self.m):
                golden[i, :] = golden[i, :] * quant_vector
        else:
            logging.info("[ERROR] can't support quant mode %s" % (self.quant_mode))

        if self.is_trans_a:
            x1_gm = x1_gm.transpose()
        if self.is_trans_b:
            x2_gm = x2_gm.transpose()
        x1_gm.tofile(work_dir + "/input/x1_gm.bin")
        x2_gm.tofile(work_dir + "/input/x2_gm.bin")
        if self.is_bias:
            bias_gm.tofile(work_dir + "/input/bias_gm.bin")
        if self.quant_mode == 2:  
            quant_vector_gm.tofile(work_dir + "/input/quant_vector_gm.bin")
        golden.tofile(work_dir + "/output/golden.bin")

        if IS_OUTPUT_TXT:
            np.savetxt(work_dir + "/input/x1_gm.txt", x1_gm.flatten(), fmt='%d', newline='\n')
            np.savetxt(work_dir + "/input/x2_gm.txt", x2_gm.flatten(), fmt='%d', newline='\n')
            np.savetxt(work_dir + "/output/golden.txt", golden.flatten(), fmt='%f', newline='\n')
            if self.is_bias:
                np.savetxt(work_dir + "/input/bias_gm.txt", bias_gm.flatten(), fmt='%d', newline='\n')
            if self.quant_mode == 2:
                np.savetxt(work_dir + "/output/quant_vector_gm.txt", quant_vector_gm.flatten(), fmt='%d', newline='\n')
        return 0


    def gen_golden_data(self, work_dir):
        if self.data_type_str == "int8_float16_dequant":
            self.gen_golden_data_int8_float16_dequant(work_dir)
        else:
            logging.info("[ERROR] can't support data type %s" % (self.data_type_str))
            return -1
        return 0
